import json
from func.gobal.data import LLmData
from openai import OpenAI
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
from func.log.default_log import DefaultLog
from func.search.search_core import SearchCore
from func.translate.baidu_translate import BaiduTranslate
import threading
import re
import time,copy

class LiteratureReviewWriter:
    llmData = LLmData()
    def __init__(self):
        self.logger = DefaultLog().getLogger()
        self.glm_client = OpenAI(api_key=self.llmData.api_key, base_url=self.llmData.base_url)
        self.baiduTranslate = BaiduTranslate()
        self.searchCore = SearchCore()

    def call_openai(self, query, system_prompt=None, model=llmData.model_name,json_format=False):
        if system_prompt is not None:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": query}
            ]
        else:
            messages = [
                {"role": "user", "content": query}
            ]
        if json_format:
            response_format={"type": "json_object"}
        else:
            response_format={"type": "text"}
        try:
            response = self.glm_client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=0.1,
                response_format=response_format
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"调用API时发生错误: {str(e)}")
            raise e  # 抛出异常供调用者捕捉

    # 判断综述类型
    def topic_classification(self, topic):
        """根据用户给出的研究主题，判断该主题属于哪一种类型的综述。"""
        query = f"""根据用户给出的研究主题，判断该主题属于哪一种类型的综述。仅返回类型即可，不要返回任何其他内容。
        
        用户给出的研究主题是：{topic}

        ### 参考样例
        Q:损失函数
        A:{{type: 对一个技术概念的调研综述}}
        Q:Text2SQL研究现状如何，面临哪些挑战？ 
        A:{{type: 对一个方向研究现状的综述}}
        Q:有哪些方法可以提升大模型的规划能力，各自优劣是什么？
        A:{{type: 对多篇文献的对比分析综述}}
        Q:多模态大模型的技术发展路线是什么样的？ 
        A:{{type: 对一个技术方法研究脉络的综述}}
        """
        response = self.call_openai(query, model=self.llmData.extract_model_name)
        return response

    # 生成文献综述大纲
    def get_structure(self, topic, topic_type, data):
        """根据用户给出的研究主题，以及文献综述的策略，生成一个文献综述大纲"""

        query = f"""
        <文献内容>
        {data}
        </文献内容>
        以上是一些与{topic}相关的文献资料。
        请分析以上文献资料，为我生成一个全面的技术调研综述文章的详细大纲，
        主题是{topic}。文章的主要研究方向是{topic_type}。
        要求大纲条理清晰，章节标题契合，前后逻辑关系或脉络关系通顺。
        每部分包含必要的子标题，尽可能考虑全面且符合逻辑的内容结构。
        帮助我生成有逻辑、有深度的技术文章大纲，不要返回任何其他内容，不要任何多余的解释。
        生成至少10个章节。
        使用英文输出。
        write in English
        """
        string = """
        ### 输出格式：
        请严格按照如下格式仅输出纯JSON格式，不要输出Python代码或其他信息，JSON字段使用顿号【、】区隔：
        {
            "title": "文章标题",
            "sections": [
                {
                    "section_title": "章节标题", 
                    "sub_sections": [
                        {
                            "sub_section_title": "子标题"
                        }
                    ]
                }
            ]
        }
        """
        structure = self.call_openai(query + string, model=self.llmData.structure_model_name, json_format=True)
        return structure

    def write_subsections_content(self, keywords, context_variables,structure, word_num=300):
        topic = context_variables["title"]
        section_title = context_variables["section_title"]
        sub_section_title = context_variables["sub_section_title"]

        main_word = keywords["main_word"]
        keyword_list = keywords["keyword_list"]
        query = f"{topic}: {section_title}-{sub_section_title}"
        #query = self.baiduTranslate.translate(query, "zh", "en")
        self.logger.info(f"提纲章节问题: {query}")
        # 本小节的参考文献搜索
        whole_content = self.get_reference_content_v2(main_word, keyword_list, topic, section_title, sub_section_title, threshold=0.55, top_k=10, sec_top_k=4)
        self.logger.info(f"{section_title}-{sub_section_title}，章节文献搜索: {len(whole_content)}")
        # for v in whole_content.values():
        #     self.logger.info(v["reference"])
        content = whole_content.values()

        SUBSECTION_WRITING_PROMPT = f'''
        你是一位人工智能领域的专家，旨在撰写一篇关于{topic}的全面综述。

        你已经创建了如下总体大纲：
        <大纲>
        {structure}
        </大纲>

        以下是参考文献列表：
        <参考文献>
        {content}
        </参考文献>

        <instruction> 
        现在你需要为以下小节撰写内容:"[{sub_section_title}]"，它属于章节："[{section_title}]"。 
        关于该小节"[{sub_section_title}]"的具体撰写要求如下描述，请务必遵循以下要求。
        
		--- [内容撰写规范] ---
        以下是撰写综述小节时引用参考文献的简要指南：
        1. 总结研究：在总结现有文献时引用相关来源。
        2. 使用具体概念或数据：讨论具体理论、模型或数据时提供引用。
        3. 比较研究发现：在比较或对比不同研究结果时引用相关研究。
        4. 指出研究空白：在指出你的综述所涉及的研究空白时引用相关文献。
        5. 使用已建立的方法：在你的综述中采用某些方法时引用其创建者。
        6. 支持论点：引用能够支持你的结论和论点的来源。
        7. 提出未来研究方向：参考与建议未来研究方向相关的研究。
        8. 你撰写的内容字数必须超过[{word_num}]。
        
        --- [引用格式规范] ---
		当文献的主要部分支持你的论点时，请严格遵循以下引用格式：
		1. 文中内引用：使用上标数字<sup>编号</sup>格式，每个引用对应唯一的编号
			- 同一句引用多篇文献时用连续上标：<sup>1</sup><sup>2</sup>
            - 如果多句引用同一篇文献，则使用同一个上标即可：<sup>1</sup>
			- 示例：...模型如<sup>17</sup><sup>18</sup>和<sup>9</sup>...
        </instruction> 

        ### 输出示例
		请严格按照如下格式仅输出JSON，不要输出Python代码或其他信息，JSON字段使用顿号【、】区隔：
		{{
			"content": "To the best of our knowledge, however, the Transformer is the first transduction model relying entirely on self-attention to compute representations of its input and output without using sequencealigned RNNs or convolution. 
            In the following sections, we will describe the Transformer, motivate self-attention and discuss its advantages over models such as <sup>1</sup><sup>2</sup> and <sup>3</sup>."
		}}

        仅返回你为小节"[{sub_section_title}]"撰写的内容(content)，字数必须超过[{word_num}]，不包含其他任何信息： 
        必须使用英文输出，不要使用中文。
        '''

        response_dict = self.call_openai(SUBSECTION_WRITING_PROMPT, model=self.llmData.model_name, json_format=True)

        if isinstance(response_dict, str):
            # 移除任何前后空白字符和控制字符
            response_dict = re.sub(r'[\x00-\x1F\x7F]', '', response_dict.strip())
            try:
                response = json.loads(response_dict)
                content = response['content']
            except Exception:
                content = str(response_dict)
                if "$your_answer" in content:
                    self.logger.exception(f"异常json:{response_dict}")
                else:
                    self.logger.error(f"异常json:{response_dict}")

        citations = list(set(re.findall(r'<sup>(\d+)</sup>', content)))
        self.logger.info(f"小节主题: {sub_section_title} 小节内容: {content} 内容引用编码: {citations}")
        reference = []

        # 创建一个新的列表来存储有效的引用
        valid_citations = []
        for citation in citations:
            if int(citation) in whole_content:
                valid_citations.append(citation)
                reference.append(whole_content[int(citation)]['reference'])
            else:
                content = re.sub(f'<sup>{citation}</sup>', '', content)

        citations = valid_citations  # 更新citations为有效的引用列表
        self.logger.info(f"替换后=> 小节主题: {sub_section_title} \n小节内容: {content},\n内容引用编码: {valid_citations},\n引用内容: {reference}")

        if len(reference) != len(valid_citations):
            self.logger.info(f"引用内容与编码不一致, 文中引用文献数: {len(valid_citations)}, 参考文献数: {len(reference)}")
            self.logger.info(f"有效引用编码: {valid_citations}")
            self.logger.info(f"参考文献: {reference}")

        final_response = {
            "content": content,
            "reference": reference
        }
        # except Exception as e:
        #     print(f"发生错误: {str(e)}")
        #     final_response = {
        #         "content": "内容生成失败",
        #         "reference": []
        #     }

        return final_response

    #注释掉原有，增加重试机制；
    def write_subsections_content_parallel_old(self, context_variables_list, structure, word_num=300):
        """并行处理多个小节的内容撰写
        
        Args:
            context_variables_list: 包含多个小节信息的列表
            structure: 文章结构
            word_num: 每个小节的最小字数要求
            
        Returns:
            results: 包含所有小节内容的字典
        """
        # 创建线程锁
        lock = threading.Lock()
        results = {}

        def process_subsection(context_variables):

            # 调用原有的write_subsections_content函数
            response = self.write_subsections_content(context_variables, structure, word_num)

            # 将结果存入字典f
            with lock:
                results[context_variables["sub_section_title"]] = response

        # 使用线程池并行处理
        with ThreadPoolExecutor(max_workers=32) as executor:
            list(tqdm(executor.map(process_subsection, context_variables_list),
                     total=len(context_variables_list),
                     desc="Writing subsections"))

        return results



    #注释掉原有，增加重试机制；
    def write_subsections_content_parallel(self, keywords, context_variables_list, structure, word_num=300, max_workers=8, delay=0.5, max_retries=5):
        """并行处理多个小节的内容撰写，但控制并发数量和增加失败后的重试
        
        Args:
            context_variables_list: 包含多个小节信息的列表
            structure: 文章结构
            word_num: 每个小节的最小字数要求
            max_workers: 最大并发线程数，默认设置为 8
            delay: 每个任务之间的延迟时间，单位秒，默认设置为 0.5秒
            max_retries: 最大重试次数，默认设置为 3
            
        Returns:
            results: 包含所有小节内容的字典
        """
        # 创建线程锁
        lock = threading.Lock()
        results = {}

        def process_subsection(context_variables):
            attempts = 0

            while attempts < max_retries:
                try:
                    # 添加时延，避免过于频繁的并发执行
                    time.sleep(delay)  # 固定的时延，可以更改为随机延迟

                    # 调用原有的write_subsections_content函数
                    response = self.write_subsections_content(keywords, context_variables, structure, word_num)

                    # 将结果存入字典
                    with lock:
                        results[context_variables["sub_section_title"]] = response
                    break  # 成功后跳出循环

                except Exception as e:
                    attempts += 1
                    self.logger.exception(f"Error occurred while processing subsection '{context_variables['sub_section_title']}': {e}")
                    if attempts < max_retries:
                        # 如果失败，等待一秒再重试
                        time.sleep(5)
                    else:
                        # 如果超过最大重试次数，记录失败
                        with lock:
                            results[context_variables["sub_section_title"]] = {
                                "content": f"Failed after {max_retries} attempts: {str(e)}",  # 失败原因直接放入 content
                                "reference": ""  # reference 为空
                            }
                    # 如果所有小节都失败，可以直接退出程序
                    self.logger.error(f"Subsection '{context_variables['sub_section_title']}' failed after {max_retries} attempts. ")
                    #sys.exit(1)  # 退出程序并返回错误状态码

        # 使用线程池并行处理，控制最大线程数
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            list(tqdm(executor.map(process_subsection, context_variables_list),
                    total=len(context_variables_list),
                    desc="Writing subsections"))

        return results



    def get_reference_content(self,query, threshold=0.6, top_k=10):
        """根据文献标题及内容，搜索参考文献"""

        # 使用字典存储,以id为key
        title_content = {}

        reference = self.searchCore.search_konwledge_content(query)
        for index, value in enumerate(reference):
            if value['distance'] > threshold:
                entity = value['entity']
                reference = self.parse_reference(value, index)
                title_content[index] = {
                    "id": index,
                    "title": entity['paper_title'],
                    "content": entity['chunk_text'],
                    "reference": reference
                }

        # 只返回前top_k个结果
        return {k:v for k,v in title_content.items() if k < top_k}

    def get_reference_content_v2(self, main_word, keyword_list, topic, section_title, sub_section_title, threshold=0.6, top_k=10, sec_top_k=2):
        """根据文献标题及内容，搜索参考文献"""

        # 使用字典存储,以id为key
        title_content = {}
        query = f"{topic}: {section_title}-{sub_section_title}"
        section = section_title + "," + sub_section_title
        keyword_list_temp = copy.deepcopy(keyword_list)
        keyword_list_temp["keyword_list"].insert(0, query)

        num = 1
        inner_num = 1
        for keyword in keyword_list_temp["keyword_list"]:
            # 首个关键字是全标题，不用组合
            if num == 1:
                last_word = keyword  #标题搜索
                reference_list = self.searchCore.search_konwledge_content(last_word, top_k)
            else:
                last_word = section + "," + main_word + "," + keyword  #关键字组合搜索
                reference_list = self.searchCore.search_konwledge_content(last_word, sec_top_k)


            for index, value in enumerate(reference_list, start=1):
                if value['distance'] > threshold:
                    id = value['id']
                    if id not in title_content:
                        entity = value['entity']
                        reference = self.parse_reference(value, inner_num)
                        title_content[id] = {
                            "id": inner_num,
                            "title": entity['paper_title'],
                            "content": entity['chunk_text'],
                            "reference": reference
                        }
                        inner_num = inner_num + 1

            num = num + 1

        # 只返回前top_k个结果
        return {index:v for index,v in enumerate(title_content.values(), start=1)}

    def parse_reference(self, data, index):
        """解析参考文献格式
        Args:
            data: 包含文献信息的字典
        Returns:
            str: 格式化的参考文献字符串
        """
        # 从文件名中提取会议名称和年份
        filename = data['entity']['original_filename']
        conference_with_year = filename.split('_')[4] # 提取会议名称(可能包含年份)

        # 尝试从会议名称中提取年份
        year_match = re.search(r'(\d{4})', conference_with_year)
        if year_match:
            # 如果会议名称中包含年份,分离会议名和年份
            year = year_match.group(1)
            conference = conference_with_year.replace(year, '')
        else:
            # 如果会议名称中没有年份,从文件名中提取
            conference = conference_with_year
            year = filename.split('_')[5].replace('.db','')

        # 组装参考文献格式
        reference = f"[{index}] {data['entity']['paper_title']}, {conference}, {year}, chunk {data['entity']['chunk_id']}"

        return reference

    def write_literature_review(self, context_variables, data):
        """根据用户给出的研究主题，以及文献综述的结构，生成一个文献综述的初稿"""
        topic = context_variables["topic"]
        structure = context_variables["structure"]
        query = f"""
        <文献资料>
        {data}
        </文献资料>
        <文献综述结构>
        {structure}
        </文献综述结构>

        ## 任务
        根据用户给出的研究主题{topic}，以及文献综述的结构\n{structure}，生成一个文献综述的初稿。
        
        ## 正文要求
        1. 文献综述的初稿应符合学术规范，逻辑清晰，内容连贯，前后逻辑关系或脉络关系通顺。
        2. 每章节包含必要的子标题，段落内容要与章节标题契合，尽可能考虑全面且符合逻辑的内容结构。
        3. 每个段落内容需要详细展开，包括但不限于:
           - 该技术/方法的基本原理和核心思想
           - 技术实现的关键步骤和算法细节
           - 与其他相关技术的对比分析
           - 该技术/方法的优势和局限性
           - 实际应用场景和效果分析
        4. 每个段落至少包含300字的详细阐述，避免笼统概括。
        5. 内容长度不低于3000字，确保每个技术点都得到充分展开。
        6. 对于每个技术要点，需要:
           - 详细说明其产生的背景和解决的问题
           - 分析其工作原理和实现机制
           - 对比不同文献中的具体实现方式
           - 总结其优缺点和适用场景

        ## 引用要求
        1. 引用文献资料时，请严格按照AMA格式插入引用，并提供完整的参考文献列表。每个关键技术点至少引用2-3篇相关文献支撑。
        2. 聚合相似引用：若多篇论文得出相似结论或使用类似方法，应将其归纳总结并同时引用。
        例如，第 2 篇和第 6 篇论文均采用相似的方法，应合并分析其异同点并一起引用。
        3. 引用格式：文章内用<sup>number</sup>，文章末尾用【标题+所在会议/期刊名称+年份+chunk序号】」的规则来进行综述撰写
        例如：
        [1] Attention is all you need, NeuralIPS, 2022, chunk 1
        [2] Layer normalization. NeuralIPS, 2022, chunk 2
        """
        response = self.call_openai(query)
        return response

    def get_strategy(self, context_variables):
        topic = context_variables["topic"]
        """根据用户给出的研究主题，生成一个文献综述的策略"""
        query = f"""请提供撰写高质量文献综述的详细策略，包括如何构建逻辑清晰的综述结构、如何对文献进行批判性分析、以及如何保证引用的全面性和学术规范性。
        请结合文献综述的目的和读者预期进行具体说明。用户给出的研究主题是：{topic}"""
        response = self.call_openai(query)
        return response

    def summarize_knowledge(self, context_variables):
        """根据用户给出的研究主题，以及文献综述的结构，生成一个文献综述总结"""
        try:
            topic = context_variables["topic"]
            data = context_variables["data"]
            query = f"""<文献内容>
            {data}
            </文献内容>
            我需要你对上述【数量】篇文献，结合主题{topic}，分别进行文献综述。
            要求给出每篇文献的简要、研究方法、研究结论、研究不足及改进建议，以及论文的优缺点。
            不要返回任何其他内容，不要任何多余的解释。

            ### 输出格式为：
            1. 文献标题
            2. 文献简要
            3. 研究方法
            3.1 方法1
            3.2 方法2
            ...
            4. 研究结论
            5. 研究不足及改进建议
            6. 论文的优缺点
            """
            response = self.call_openai(query)
            return response
        except Exception as e:
            print(f"生成文献综述总结时发生错误: {str(e)}")
            return ""

    def recursive_summarize(self, context_variables, k=5):
        """递归地对文本块进行分组总结
        
        Args:
            text_blocks: 待总结的文本块列表
            k: 每组的文本块数量
            
        Returns:
            最终总结的知识点列表
        """
        text_blocks = context_variables["data"]
        if len(text_blocks) <= k:
            return text_blocks

        # 按k个一组分割文本块
        grouped_blocks = [text_blocks[i:i+k] for i in range(0, len(text_blocks), k)]
        summarized_blocks = []

        # 对每组进行总结
        for group in grouped_blocks:
            group_text = "\n==========================\n".join(group)
            topic = context_variables.get("topic", "")
            summary = self.summarize_knowledge({"topic": topic, "data": group_text})
            summarized_blocks.append(summary)

        # 如果总结后的块数仍然大于k,继续递归总结
        if len(summarized_blocks) > k:
            return self.recursive_summarize(summarized_blocks, k)

        return summarized_blocks

    def process_item(self, item, topic):
        """处理单个文献项
        
        Args:
            item: 文献项
            topic: 研究主题
        
        Returns:
            summary: 总结内容
        """
        paper_text = item['chunk_text']
        context_variables = {"topic": topic, "data": paper_text}
        summary = self.summarize_knowledge(context_variables)
        return summary

    def process_data_parallel(self, merged_data, topic):
        """并行处理数据
        
        Args:
            merged_data: 需要处理的数据字典
            topic: 研究主题
            
        Returns:
            处理后的数据字典
        """
        # 创建线程锁
        lock = threading.Lock()

        def process_with_progress(item_id, item):
            summary = self.process_item(item, topic)
            with lock:
                merged_data[item_id]['summary'] = summary

        # 使用线程池并行处理
        with ThreadPoolExecutor(max_workers=5) as executor:
            list(tqdm(executor.map(lambda x: process_with_progress(x[0], x[1]),
                                 merged_data.items()),
                     total=len(merged_data)))

        return merged_data

    def retrieve_prompt(self, topic, language):
        prompt = ""
        if language == "中文":
            prompt = f"""我现在需要撰写一篇关于{topic}的文献综述。请启动联网功能，寻找2020年之后引用次数最多的10篇相关文献。
            你的回复格式应包括：标题、作者、发表年份、引用次数，以及查看链接。"""
        else:
            prompt = f"""I need to write a literature review on the topic of {topic}. 
            Please enable the retrieve function and find the top 10 most cited papers published after 2020. 
            The response should include: title, author, publication year, citation count, and link to view."""
        return prompt

    def summarize_prompt(self, texts, language):
        prompt = ""
        if language == "中文":
            prompt = f"""请对上述找到的【数量】篇文献分别进行综述，内容包括：研究背景与目的、研究方法、研究结论、研究不足及改进建议，以及论文的优缺点。
            请严格按照AMA格式插入引用，并提供完整的参考文献列表。"""
        else:
            prompt = f"""Please provide a review of each of the [Number] papers found above, including: research background and objectives, research methods, research conclusions, limitations and suggestions for improvement, and the strengths and weaknesses of the paper.{texts}"""
        return prompt

    def put_togeter_prompt(self, texts, language):
        prompt = ""
        if language == "中文":
            prompt = f"""将上述【数量】篇文献的简要综述整合为一个符合学术规范的文献综述段落，确保逻辑清晰、内容连贯，并在适当位置插入引用文献。严格按照AMA格式插入引用，并提供完整的参考文献列表。{texts}"""
        else:
            prompt = f"""Integrate the summaries of the above [Number] papers into a cohesive literature review paragraph that adheres to academic standards. Ensure clear logic and coherence, and insert citations in appropriate places. Strictly follow the AMA format for citations and provide a complete reference list.{texts}"""
        return prompt